% Copyright (C) 2019-2023 Benjamin Born, Francesco D'Ascanio, Gernot J. Mueller, Johannes Pfeifer
%
% This is free software: you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
%
% It is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
% 
% For a copy of the GNU General Public License,
% see <http://www.gnu.org/licenses/>.

% Compute impulse responses to positive and negative g shocks under full
% employment

clear;
if ~isfolder('Figures')
    mkdir('.','Figures');
end
if ~isfolder('IRF_results')
    mkdir('.','IRF_results');
end

addpath('Auxiliary_functions/')

IRF_save_name = 'IRF_results/IRF_peg' ;

N_periods = 150; %number of simulation periods (quarters)
N_rep = 1e6; %number of replications
T_ergodic = 1e6;1e8; %periods for simulation of mean;
T_burnin = 1e+4; %burn-in period

%load policy function
loaded_PFI=load('Policy_functions/pfi_peg_g_low_beta_large_grid.mat');

if ~isfield(loaded_PFI.grid,'w_grid_vec')
    loaded_PFI.grid.w_grid_vec=loaded_PFI.w_grid_vec;
end

omega=loaded_PFI.par.omega;
alfa=loaded_PFI.par.alfa;
hbar=loaded_PFI.par.hbar;
xi=loaded_PFI.par.xi;
if ~isfield(loaded_PFI.par,'phi_e')
    loaded_PFI.par.phi_e= 0; %exchange rate policy parameter
end


%% Set initial conditions of exogenous states, debt, and past wage equal to their respective unconditional means
temp=unique(loaded_PFI.grid.y_level);
y_0 = temp(round(length(temp)/2));
temp=unique(loaded_PFI.grid.r_level);
r_0 = temp(round(length(temp)/2));
temp=unique(loaded_PFI.grid.g_level);
g_0 = temp(round(length(temp)/2));
g_0_pos = loaded_PFI.grid.g_level(find(loaded_PFI.grid.g_level==g_0,1,'last')+1);
g_0_neg = loaded_PFI.grid.g_level(find(loaded_PFI.grid.g_level==g_0, 1 )-1);
d_0 = loaded_PFI.grid.d_grid(315);
w_0 = loaded_PFI.grid.w_grid_vec(225);
h_0 = loaded_PFI.h(round(loaded_PFI.grid.n_y/2),315,225);

%store initial values
y_T_start = repmat(y_0,[1 3]);
r_start = repmat(r_0,[1 3]);
g_N_start = [g_0 g_0_pos g_0_neg];
d_start = repmat(d_0,[1 3]);
w_lag_start = repmat(w_0,[1 3]);
h_start = repmat(h_0,[1 3]);

c_T = y_0 + d_0./(1+r_0) - d_0;
y_N = h_0.^alfa; %nontradable output
c_N = y_0 - g_0; %nontradable consumption
p_lag_start = (1-omega)/omega*(c_N./c_T).^(-1/xi); %relative price of nontradables in terms of tradables
p_lag_start = repmat(p_lag_start,[1 3]);
clear('c_T','y_N','c_N','r_0','g_0','d_0','g_0_pos','g_0_neg','h_0');

%P_trans is the transition probability matrix of tradable output. Cumul_P_trans is the cumulative probability matrix (useful for drawing tradable output realizations)
Cum_prob = cumsum(loaded_PFI.MC.P_trans,2);

[Y_base, var_name, pos]=simulate_model_general(y_T_start(1),r_start(1),g_N_start(1),d_start(1),w_lag_start(1),p_lag_start(1),T_ergodic,loaded_PFI.dp,loaded_PFI.h,loaded_PFI.grid,loaded_PFI.par,Cum_prob,1);
Y_base_mean=mean(Y_base(T_burnin+1:end,:));

nvar = length(fieldnames(pos)); %number of variables for which IR are computed

%Initialize vector of impulse responses
Y = zeros(N_periods,nvar,N_rep);
Y_pos = zeros(N_periods,nvar,N_rep);
Y_neg = zeros(N_periods,nvar,N_rep);

%reset random number generator
rng('default')
tic
for replic_iter=1:N_rep
    %set initial condition for current replication
    y_T = y_T_start;
    r = r_start;
    g_N = g_N_start;
    d = d_start;
    w_lag= w_lag_start;
    p_lag= p_lag_start;
    H = h_start;
    Y(:,:,replic_iter)=simulate_model_general(y_T(1),r(1),g_N(1),d(1),w_lag(1),p_lag(1),N_periods,loaded_PFI.dp,loaded_PFI.h,loaded_PFI.grid,loaded_PFI.par,Cum_prob);
    Y_pos(:,:,replic_iter)=simulate_model_general(y_T(2),r(2),g_N(2),d(2),w_lag(2),p_lag(2),N_periods,loaded_PFI.dp,loaded_PFI.h,loaded_PFI.grid,loaded_PFI.par,Cum_prob);
    Y_neg(:,:,replic_iter)=simulate_model_general(y_T(3),r(3),g_N(3),d(3),w_lag(3),p_lag(3),N_periods,loaded_PFI.dp,loaded_PFI.h,loaded_PFI.grid,loaded_PFI.par,Cum_prob);    
end
toc


%% Plot IRFs
IRFs=NaN(nvar,N_periods,2);
h1=figure('Name','Positive shock');
for var_iter=1:min(nvar,16)
    subplot(4,4,var_iter)
    IRFs(var_iter,:,1)=mean(squeeze((Y_pos(:,var_iter,:)-Y(:,var_iter,:))),2);
    plot(IRFs(var_iter,:,1)./Y_base_mean(var_iter))
    xlim([1 N_periods]);
    title(var_name{var_iter}, 'interpreter','latex')
end
print([IRF_save_name '_positive'],'-depsc2')
saveas(h1,[IRF_save_name '_positive'])

if nvar>16
    h3=figure('Name','Positive shock 2');
    for var_iter=1:min(nvar-16,16)
        subplot(4,4,var_iter)
        IRFs(16+var_iter,:,1)=mean(squeeze((Y_pos(:,16+var_iter,:)-Y(:,16+var_iter,:))),2);
        plot(IRFs(16+var_iter,:,1)./Y_base_mean(16+var_iter))
        xlim([1 N_periods]);
        title(var_name{16+var_iter}, 'interpreter','latex')
    end
end
print([IRF_save_name '_positive_2'],'-depsc2')
saveas(h3,[IRF_save_name '_positive_2'])

h2=figure('Name','Negative shock');
for var_iter=1:min(nvar,16)
    subplot(4,4,var_iter)
    IRFs(var_iter,:,2)=mean(squeeze((Y_neg(:,var_iter,:)-Y(:,var_iter,:))),2);
    plot(IRFs(var_iter,:,2)./Y_base_mean(var_iter))
    xlim([1 N_periods]);
    title(var_name{var_iter}, 'interpreter','latex')
end
print([IRF_save_name '_negative'],'-depsc2')
saveas(h2,[IRF_save_name '_negative'])

if nvar>16
    h4=figure('Name','Negative shock 2');
    for var_iter=1:min(nvar-16,16)
        subplot(4,4,var_iter)
        IRFs(16+var_iter,:,2)=mean(squeeze((Y_neg(:,16+var_iter,:)-Y(:,16+var_iter,:))),2);
        plot(IRFs(16+var_iter,:,2)./Y_base_mean(16+var_iter))
        xlim([1 N_periods]);
        title(var_name{16+var_iter}, 'interpreter','latex')
    end    
end
print([IRF_save_name '_negative_2'],'-depsc2')
saveas(h4,[IRF_save_name '_negative_2'])

save('IRF_results/irfs_peg.mat','IRFs','Y_base_mean','N_periods','var_name','pos','-v7.3')